#encoding=utf-8
from __future__ import print_function
import sys
reload(sys)
sys.setdefaultencoding('utf-8')
import mxnet as mx
import numpy as np
import mxnet.ndarray as nd
import mxnet.gluon as gluon
import mxnet.gluon.nn as nn
from models.defomableconvnets.config import update_config,config
from models.defomableconvnets.resnet_v1_101_fpn_dcn_rcnn import resnet_v1_101_fpn_dcn_rcnn
from models.defomableconvnets.image import transform
from models.defomableconvnets.bbox.bbox_transform import bbox_pred,bbox_transform,clip_boxes
from models.defomableconvnets.nms.nms import gpu_nms_wrapper
from utils.block import SymbolBlock
from utils.common import lsdir
import cv2
def get_symbol(cfg):
    sym = resnet_v1_101_fpn_dcn_rcnn().get_symbol(cfg,is_train=False)
    return sym

def im_detect_bbox_aug(net, nms_wrapper, img_path, scales, pixel_means, ctx, threshold = 1e-3, viz = False):

    all_bboxes = []
    all_scores = []
    img_ori = cv2.imread(img_path.encode("utf-8"))
    for scale_min, scale_max in scales:
        fscale = 1.0 * scale_min / min(img_ori.shape[:2])
        img_resized = cv2.resize(img_ori,(0,0),fx = fscale,fy = fscale)
        h,w,c = img_resized.shape
        h_padded = h if h %32==0 else h + 32 - h % 32
        w_padded = w if w %32==0 else w + 32 - w % 32
        img_padded = np.zeros(shape = (h_padded,w_padded,c), dtype = img_resized.dtype)
        img_padded[:h,:w,:] = img_resized
        img = transform(img_padded,pixel_means=pixel_means)
        im_info = nd.array([[h_padded,w_padded,1.0]], ctx= ctx[0])
        data = nd.array(img,ctx=ctx[0])

        rois, scores, bbox_deltas = net(data,im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])
        bbox /= fscale
        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

        # hflip
        rois, scores, bbox_deltas = net(data[:,:,:,::-1],im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])

        tmp = bbox[:, 0::4].copy()
        bbox[:, 0::4] =data.shape[3] - bbox[:, 2::4] -1 #x0 = w - x0
        bbox[:, 2::4] =data.shape[3] - tmp -1 # x1 = w -x1
        bbox /= fscale

        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

        #
        # vflip
        rois, scores, bbox_deltas = net(data[:,:,::-1],im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])

        tmp = bbox[:, 1::4].copy()
        bbox[:, 1::4] =data.shape[2] - bbox[:, 3::4] -1 #x0 = w - x0
        bbox[:, 3::4] =data.shape[2] - tmp -1 # x1 = w -x1
        bbox /= fscale

        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())


        # vhflip
        rois, scores, bbox_deltas = net(data[:,:,::-1, ::-1],im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])

        tmp = bbox[:, 1::4].copy()
        bbox[:, 1::4] =data.shape[2] - bbox[:, 3::4]-1 #x0 = w - x0
        bbox[:, 3::4] =data.shape[2] - tmp - 1 # x1 = w -x1

        tmp = bbox[:, 0::4].copy()
        bbox[:, 0::4] =data.shape[3] - bbox[:, 2::4]-1 #x0 = w - x0
        bbox[:, 2::4] =data.shape[3] - tmp -1 # x1 = w -x1

        bbox /= fscale

        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

    all_bboxes = np.concatenate(all_bboxes,axis = 0)
    all_scores = np.concatenate(all_scores,axis = 0)
    pred_bboxes = []
    pred_scores = []
    pred_clsid = []
    for j in range(1, all_scores.shape[1]):
        cls_scores = all_scores[:, j, np.newaxis]
        cls_boxes = all_bboxes[:, 4:8] if config.CLASS_AGNOSTIC else all_bboxes[:, j * 4:(j + 1) * 4]
        cls_dets = np.hstack((cls_boxes, cls_scores))
        keep = nms_wrapper(cls_dets.astype('f'))
        cls_dets = cls_dets[keep, :]
        cls_dets = cls_dets[cls_dets[:, -1] > threshold, :]
        pred_bboxes.append(cls_dets[:,:4])
        pred_scores.append(cls_dets[:,4])
        pred_clsid.append(j * np.ones(shape=(cls_dets.shape[0],), dtype=np.int))
    pred_bboxes = np.concatenate(pred_bboxes,axis = 0)
    pred_scores = np.concatenate(pred_scores,axis = 0)
    pred_clsid = np.concatenate(pred_clsid,axis = 0)
    if viz:
        import gluoncv
        import matplotlib.pyplot as plt
        gluoncv.utils.viz.plot_bbox(img_ori[:,:,::-1], bboxes=pred_bboxes,scores=pred_scores,labels=pred_clsid,thresh=.5)
        # plt.show()
    return pred_bboxes,pred_scores,pred_clsid

def validate(net, num_wrapper, ctx_list):
    # Get validation dataset
    TEST_SCALES = [[960, 1280]]
    from data.bbox.aluminum_material import AluminumDet
    from mAP import VOCMApMetric
    val_dataset = AluminumDet(is_train=False)
    metric = VOCMApMetric(iou_thresh=.5, class_names=[str(x) for x in range(len(val_dataset.classes))])
    for i in tqdm.tqdm(range(len(val_dataset))):
        im_name, gt_boxes = val_dataset.at_with_image_path(i)
        bboxes, scores, labels = im_detect_bbox_aug(net, nms_wrapper, im_name, TEST_SCALES, config.network.PIXEL_MEANS,
                                                    ctx=ctx, viz=False)
        gt_labels = gt_boxes[np.newaxis,:,4]
        metric.update(pred_bboxes=bboxes[np.newaxis], pred_labels = labels[np.newaxis] -1 ,
                      pred_scores = scores[np.newaxis], gt_bboxes = gt_boxes[np.newaxis,:,:4],
                      gt_labels = gt_labels, gt_difficults=np.zeros_like(gt_labels))
    classes_mAP = dict(zip(*metric.get()))
    classes_mAPChinese = {}

    for k in classes_mAP:
        v = classes_mAP[k]
        if k == "mAP":
            classes_mAPChinese[k] = v
        else:
            classes_mAPChinese[val_dataset.classes[int(k)]] = v
    for k in classes_mAPChinese.keys():
        print (k, classes_mAPChinese[k])
    return classes_mAPChinese["mAP"]
if __name__ == '__main__':
    import os,tqdm,time
    os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
    gpu_id = 4
    ctx = [mx.gpu(gpu_id)]
    update_config("/data1/zyx/yks/Deformable-ConvNets/experiments/fpn/cfgs/resnet_v1_101_coco_trainval_fpn_dcn_end2end_ohem.yaml")
    sym = get_symbol(config)
    net = SymbolBlock(sym= sym,input_names=["data","im_info"],
                      resume="/data1/zyx/yks/Deformable-ConvNets/fpn/output/fpn/coco/resnet_v1_101_coco_trainval_fpn_dcn_end2end_ohem/train2014/fpn_coco-0005.params")
    net.collect_params().reset_ctx(ctx)
    im_names = list(lsdir("/data1/zyx/yks/dataset/guangdong_round2_test_a_20181011",suffix=".jpg"))
    from random import shuffle
    # shuffle(im_names)

    nms_wrapper = nms = gpu_nms_wrapper(config.TEST.NMS, gpu_id)
    #score = validate(net, nms_wrapper, ctx_list = ctx)
    results = {}
    results["results"]=[]
    for im_name in tqdm.tqdm(im_names):
        TEST_SCALES=[[960,1280]]
        one_img = {}
        one_img["filename"] = os.path.basename(im_name)
        one_img["rects"] = []
        bboxes, scores, labels = im_detect_bbox_aug(net,nms_wrapper,im_name, TEST_SCALES, config.network.PIXEL_MEANS,ctx=ctx,viz=False)
        import matplotlib.pyplot as plt
        plt.show()
        for bbox, score, label in zip(bboxes,scores,labels):
            one_rect = {}
            xmin,ymin,xmax,ymax = bbox[:4]
            one_rect["xmin"] = int(np.round(xmin))
            one_rect["ymin"] = int(np.round(ymin))
            one_rect["xmax"] = int(np.round(xmax))
            one_rect["ymax"] = int(np.round(ymax))
            one_rect["confidence"] = score
            one_rect["label"] = "defect%d"%(label-1,)
            one_img["rects"].append(one_rect)
        results["results"].append(one_img)
    import json
    json.dump(results,open("submit_testA_aluminum_minus1_vhflip_epoch5.json","wt"))
